package org.dreamwork.sso.saml2.servlet;

import org.dreamwork.sso.saml2.SamlKeys;
import org.dreamwork.sso.saml2.protocol.SamlAssertion;
import org.dreamwork.sso.saml2.protocol.SamlLoginResponse;
import org.dreamwork.sso.saml2.protocol.SamlStatus;
import org.dreamwork.sso.saml2.spec.IKeyFetcher;
import org.dreamwork.secure.AlgorithmMapping;
import org.dreamwork.sso.saml2.util.ConfigKey;
import org.dreamwork.sso.saml2.util.DynamicFromBuilder;
import org.dreamwork.misc.XMLUtil;
import org.dreamwork.gson.GsonHelper;
import org.dreamwork.misc.Base64;
import org.dreamwork.util.StringUtil;
import org.opensaml.Configuration;
import org.opensaml.xml.io.Unmarshaller;
import org.opensaml.xml.io.UnmarshallingException;
import org.opensaml.xml.signature.KeyInfo;
import org.opensaml.xml.signature.X509Certificate;
import org.opensaml.xml.signature.X509Data;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.xml.sax.SAXException;

import javax.security.cert.CertificateException;
import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import javax.xml.parsers.ParserConfigurationException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.net.URLEncoder;
import java.security.*;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * Created with IntelliJ IDEA.
 * User: seth.yang
 * Date: 13-8-26
 * Time: 下午6:12
 */
public class SamlLoginResponseHandler extends HttpServlet {
    private IKeyFetcher fetcher;
    private Class<?> entityType;
    private String defaultRedirectURL;

    private static final String ENC = "utf-8";

    @Override
    public void init () throws ServletException {
        ServletContext context = getServletContext ();
        Properties props = (Properties) context.getAttribute (SamlKeys.CONFIG_CONTEXT_KEY);
        fetcher = (IKeyFetcher) context.getAttribute (SamlKeys.KEY_FETCHER_CONTEXT_KEY);
        defaultRedirectURL = props.getProperty (ConfigKey.SAML2_SP_DEFAULT_REDIRECT_URL);

        try {
            entityType = Class.forName (props.getProperty ("saml2.sp.entity.class"));
        } catch (ClassNotFoundException ex) {
            throw new ServletException (ex);
        }
    }

    @Override
    protected void doPost (HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
        String raw = request.getParameter ("SAMLResponse");
        try {
            SamlLoginResponse samlLoginResponse = new SamlLoginResponse (raw);

            verifyResponse (request, response);

            if (samlLoginResponse.isSigned ()) {
                if (!samlLoginResponse.verifySignature ()) {
                    throw new ServletException ("SAML2 Login Response Signature Verify fail!");
                }
            }

            SamlAssertion assertion;
            if (samlLoginResponse.isAssertionEncrypted ()) {
                PrivateKey privateKey = fetcher.getPrivateKey (getServletContext (), null);
                assertion = samlLoginResponse.decryptAssertion (privateKey);
            } else
                assertion = samlLoginResponse.getAssertion ();

            Object entity = assertion.getEntity (entityType);

            System.out.println (GsonHelper.getGson ().toJson (entity));

            HttpSession session = request.getSession ();
            session.setAttribute (SamlKeys.SAML_LOGIN_STATUS, SamlStatus.SUCCESSFUL);
            session.setAttribute (SamlKeys.SAML_ENTITY_SESSION_KEY, entity);

            restorePreviousRequest (request, response);
        } catch (Exception ex) {
            throw new ServletException (ex);
        }
    }

    @SuppressWarnings ("unchecked")
    private void restorePreviousRequest (HttpServletRequest request, HttpServletResponse response) throws IOException {
        HttpSession session = request.getSession ();
        Map<String, Object> backup = (Map<String, Object>) session.getAttribute (SamlKeys.SAML_REQUEST_BACKUP);
        if (backup == null || backup.size () == 0) {
            response.sendRedirect (defaultRedirectURL);
            return;
        }

        String method = (String) backup.get ("method");
        if ("GET".equalsIgnoreCase (method)) {
            StringBuffer url = (StringBuffer) backup.get ("url");
            String queryString = (String) backup.get ("queryString");
            if (!StringUtil.isEmpty (queryString))
                url.append ('?').append (queryString);
            response.sendRedirect (url.toString ());
        } else if ("POST".equalsIgnoreCase (method)) {
            StringBuffer action = (StringBuffer) backup.get ("url");
            Map<String, Object> parameters = (Map<String, Object>) backup.get ("parameters");
            DynamicFromBuilder builder = new DynamicFromBuilder (action.toString (), parameters);
            builder.build (response.getWriter ());
            response.flushBuffer ();
        }
    }

    private void verifyResponse (HttpServletRequest request, HttpServletResponse response) throws IOException, ParserConfigurationException, SAXException, UnmarshallingException, ServletException, CertificateException, NoSuchAlgorithmException, InvalidKeyException, java.security.cert.CertificateException, SignatureException {
        String algorithm = request.getParameter ("SigAlg");
        String signedText = request.getParameter ("Signature");
        String keyInfo = request.getParameter ("KeyInfo");
        String samlResponseText = request.getParameter ("SAMLResponse");

        if (StringUtil.isEmpty (algorithm)) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Unknown signature algorithm");
            return;
        }

        if (StringUtil.isEmpty (keyInfo)) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Can't find key info");
            return;
        }

        ByteArrayInputStream in = new ByteArrayInputStream (Base64.decode (keyInfo));
        Document doc = XMLUtil.parse (in);
        Element root = doc.getDocumentElement ();

        Unmarshaller unmarshaller = Configuration.getUnmarshallerFactory ().getUnmarshaller (root);
        KeyInfo ki = (KeyInfo) unmarshaller.unmarshall (root);
        List<X509Data> list = ki.getX509Datas ();
        if (list == null || list.size () == 0) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Can't find x509 certificate in KeyInfo");
            return;
        }

        X509Data data = list.get (0);
        List<X509Certificate> certs = data.getX509Certificates ();
        if (certs == null || certs.size () == 0) {
            response.sendError (HttpServletResponse.SC_BAD_REQUEST, "Can't find x509 certificate in KeyInfo");
            return;
        }

        X509Certificate cert = certs.get (0);
        String text = cert.getValue ();
        byte[] certData = Base64.decode (text);

        String s_response = URLEncoder.encode (new String (Base64.decode (samlResponseText), ENC), ENC);
        String s_algorithm = URLEncoder.encode (algorithm, ENC);
        String rawText = "SAMLResponse=" + s_response + "&SigAlg=" + s_algorithm;

        java.security.cert.CertificateFactory cf = java.security.cert.CertificateFactory.getInstance ("X.509");
        java.security.cert.Certificate x509 = cf.generateCertificate (new ByteArrayInputStream (certData));

        String jceName = AlgorithmMapping.translateSignatureAlgorithm (algorithm);
        Signature signature = Signature.getInstance (jceName);
        signature.initVerify (x509.getPublicKey ());
        signature.update (rawText.getBytes (ENC));
        if (!signature.verify (Base64.decode (signedText))) {
            response.sendError (HttpServletResponse.SC_UNAUTHORIZED, "Can't find x509 certificate in KeyInfo");
        }
    }
}